import cv2
import os
import subprocess
import time
import numpy as np
from threading import Lock, Event, Thread
from vehicle_flow_counting import CONFIG_FILE, V_OUT_FILE, DISPLAY_PUSH
from .config_loader import load_config
from .frame_display import show_frame_scaled

class VideoIO:
    """
    视频读写封装类
    功能：
    - 打开视频源（文件或摄像头）
    - 获取视频参数（帧率、宽高）
    - 初始化视频写入对象
    - 自动释放资源
    """
    def __init__(self, config_path=CONFIG_FILE, output_file=V_OUT_FILE, codec='mp4v'):
        """
        参数：
            input_url (str): 视频源路径或摄像头 URL
            output_file (str, optional): 输出视频路径，不提供则不保存
            codec (str): 本地保存视频的编码方式，如 'mp4v', 'XVID', 'MJPG' (默认 mp4v)
            resize (tuple, optional): (width, height)，如果指定会调整输出分辨率
        """
        # 读取配置
        self.command, self.input_url, self.classes, self.out_width, self.out_height ,out_fps = load_config(config_path)
        self.output_file = output_file
        self.push_fps = out_fps  # 推流帧率（None 表示无控制）

        # 打开视频源
        self.cap = cv2.VideoCapture(self.input_url)
        if not self.cap.isOpened():
            raise RuntimeError(f"无法打开视频源: {self.input_url}")

        # 获取视频参数
        self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))

        # 初始化视频写入对象
        self.writer = None
        if output_file:
            fourcc = cv2.VideoWriter_fourcc(*codec) # 保存视频的编码方式 可换成 'mp4v'
            if os.path.exists(output_file):
                os.remove(output_file)  # 避免 FileExistsError
            self.writer = cv2.VideoWriter(output_file, fourcc, self.fps, (self.width, self.height))

        # 异步推流相关
        # 初始化 ffmpeg 推流管道
        # 创建子进程 启动 FFmpeg 命令 让它在后台运行；定义子进程的标准输入（stdin）被重定向为管道；通过 p.stdin.write() 直接把数据传给 FFmpeg
        if DISPLAY_PUSH:
            self.ffmpeg_process = subprocess.Popen(self.command, stdin=subprocess.PIPE, stdout=None, stderr=None, bufsize=10**6 ) # 100MB 缓冲区
            time.sleep(0.1)  # 等待 FFmpeg 初始化
            self.stop_flag = Event()

            # 新增：存放最新帧 + 锁
            self.latest_frame = None
            self.latest_frame_lock = Lock()

            # 启动推流线程
            self.thread = Thread(target=self._stream_worker, daemon=True)
            self.thread.start()


    def get_params(self):
        """
        获取当前视频参数
        """
        return self.width, self.height, self.fps, self.input_url, self.classes

    def read_frame(self):
        """
        读取一帧视频，如果 resize 会调整分辨率
        返回：
            frame (np.array) 或 None
        """
        ret, frame = self.cap.read()
        if not ret:
            return None
        return frame

    def write_frame(self, frame):
        """
        写入一帧到输出视频（如果指定了 writer）
        """
        if self.writer is not None:
            self.writer.write(frame)

    def resize_frame(self, frame):
        # 调整分辨率（仅用于推流）
        # print(frame.shape[1], frame.shape[0])
        if (frame.shape[1], frame.shape[0]) != (self.out_width, self.out_height):
            frame = cv2.resize(frame, (self.out_width, self.out_height))
        return frame

    # def send_frame(self, frame):
    #     # 调整分辨率为配置文件中的
    #     frame_resized = self.resize_frame(frame)
    #     # 写入 FFmpeg stdin
    #     if self.ffmpeg_process:
    #         self.ffmpeg_process.stdin.write(frame_resized.tobytes())

    def _start_ffmpeg_process(self):
        """
        启动或重启 FFmpeg 推流进程
        """
        if hasattr(self, 'ffmpeg_process') and self.ffmpeg_process is not None:
            try:
                self.ffmpeg_process.stdin.close()
            except Exception:
                pass
            try:
                self.ffmpeg_process.kill()
            except Exception:
                pass
            self.ffmpeg_process = None

        print("[VideoIO] Starting FFmpeg process...")
        self.ffmpeg_process = subprocess.Popen(
            self.command,                   # 之前 load_config 返回的命令列表
            stdin=subprocess.PIPE,          # 接收帧数据
            # stdout=subprocess.DEVNULL,      # 不显示 FFmpeg 输出
            # stderr=subprocess.PIPE,         # 错误信息可调试
            bufsize=10**6                   # 大缓冲区防止卡顿
        )
        time.sleep(0.1)  # 等待 FFmpeg 初始化


    def _stream_worker(self):
        interval = 1 / self.push_fps if self.push_fps else 0
        print(f"interval1: {1000*interval} ms")
        last_frame = np.zeros((self.out_height, self.out_width, 3), dtype=np.uint8)
        frame_id_ = 0

        while not self.stop_flag.is_set():
            start_time = time.time()

            # 取最新帧
            with self.latest_frame_lock:
                if self.latest_frame is not None:
                    frame_to_send = self.latest_frame.copy()
                    last_frame = frame_to_send
                else:
                    frame_to_send = last_frame  # 如果还没有帧，就复用上一帧

            # 检查 FFmpeg 进程是否存活
            if self.ffmpeg_process is None or self.ffmpeg_process.poll() is not None:
                print("[Stream Worker] FFmpeg process not running, restarting...")
                self._start_ffmpeg_process()  # 你需要实现启动 FFmpeg 的函数
                time.sleep(0.1)  # 等待 FFmpeg 初始化

            # 送进 FFmpeg
            if self.ffmpeg_process:
                try:
                    frame_id_ += 1
                    self.ffmpeg_process.stdin.write(frame_to_send.tobytes())
                    print(" ")
                    print("frame_id_: ", frame_id_)
                except BrokenPipeError:
                    print("[Stream Worker] Broken pipe detected, will restart FFmpeg.")
                    if self.ffmpeg_process:
                        try:
                            self.ffmpeg_process.stdin.close()
                        except Exception:
                            pass
                        self.ffmpeg_process.kill()
                    self.ffmpeg_process = None
                    continue  # 下一轮循环自动重启 FFmpeg

            # 控制输出 fps
            if self.push_fps > 0:
                elapsed = time.time() - start_time
                sleep_time = max(0, interval - elapsed)
                print(f"elapsed: {1000*elapsed} ms")
                print(f"interval2: {1000*interval} ms")
                print(f"sleep_time: {1000*sleep_time} ms")
                # print("Thread end\n")
                time.sleep(sleep_time)

    def push_frame(self, frame):
        frame_resized = self.resize_frame(frame)
        # print("latest_frame_lock:", self.latest_frame_lock)
        with self.latest_frame_lock:
            self.latest_frame = frame_resized

    def stop_push(self):
        self.stop_flag.set()
        self.thread.join(timeout=1.0)
        if self.ffmpeg_process:
            try:
                self.ffmpeg_process.stdin.close()
                self.ffmpeg_process.wait(timeout=2)
            except Exception:
                self.ffmpeg_process.kill()

    def play_input_video(self, path):
        """
        播放视频并支持逐帧处理
        Args:
            path (str): 视频文件路径 或 摄像头ID(0/1/2...)
        """
        cap = cv2.VideoCapture(path)
        if not cap.isOpened():
            raise RuntimeError(f"无法打开视频源: {path}")

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            exit_flag = show_frame_scaled(frame, window_name="Input video stream", mode="video", scale_factor=0.7, show=True)

            if exit_flag:  # 空格键退出
                break

        cap.release()
        cv2.destroyAllWindows()

    def release(self):
        """
        释放资源
        """
        if self.cap:
            self.cap.release() # 释放打开的（原）视频对象
        if self.writer:
            self.writer.release() # 释放视频写入对象
        if DISPLAY_PUSH:
            if self.ffmpeg_process:
                self.ffmpeg_process.stdin.close() # 释放推流管道对象
                self.ffmpeg_process.wait()
            self.stop_push()
        cv2.destroyAllWindows() # 关闭所有窗口

        # 针对writer(会默认创建一个空文件)：检查输出文件是否小于 10kB，如果是就删除
        if self.output_file and os.path.exists(self.output_file):
            # print(os.path.getsize(self.output_file))
            if os.path.getsize(self.output_file) < 10240:
                os.remove(self.output_file)
                # print(f"输出文件过小(<10kB)，已删除: {self.output_file}")
